import torch
from torch.utils.data import Dataset
import numpy as np
import pandas as pd
from PIL import Image
import os
import re
import torchvision
from torch import nn
from torchvision import models
from torchvision import transforms

""" 
class Dataset(Dataset):
    def __init__(self, data_label_list_path, data_vgg_list_path, data_bert_list_path, index_id=None):
        data_labels = np.array(torch.load(data_label_list_path)).astype("int")
        data_vgg = torch.load(data_vgg_list_path)
        x_features = torch.load(data_bert_list_path)


        if index_id != None:
            data_labels = data_labels[index_id]
            data_vgg = data_vgg[index_id]
            x_features = x_features[index_id]

        self.data_labels = data_labels
        self.data_vgg = data_vgg
        self.x_features = x_features

        self.tokenizer = BertTokenizer.from_pretrained('/root/autodl-tmp/bert-base-uncased', do_lower_case=True)



    def __len__(self):
        return len(self.data_labels)

    def __getitem__(self, index):
        tensor_input_id, tensor_input_mask = self.pre_processing_BERT(text)
        return index, \
               self.data_labels[index], \
               self.data_vgg[index], \
               self.x_features[index]
"""
def text_preprocessing(text):
    """
    - 删除实体@符号(如。“@united”)
    — 纠正错误(如:'&amp;' '&')
    @参数 text (str):要处理的字符串
    @返回 text (Str):已处理的字符串
    """
    # 去除 '@name'
    text = re.sub(r'(@.*?)[\s]', ' ', text)

    #  替换'&amp;'成'&'
    text = re.sub(r'&amp;', '&', text)

    # 删除尾随空格
    text = re.sub(r'\s+', ' ', text).strip()

    return text

class Dataset(Dataset):

    def __init__(self, df_path, root_dir, tokenizer, MAX_LEN, div_data = 1):
        """
        参数:
            csv_file (string):包含文本和图像名称的csv文件的路径
            root_dir (string):目录
            transform(可选):图像变换
        """
        self.csv_data =  pd.read_csv(df_path, sep='\t', header=0, index_col='index')
        self.csv_data = self.csv_data[:int(self.csv_data.shape[0]/div_data)]

        self.root_dir = root_dir
        self.tokenizer_bert = tokenizer

        self.MAX_LEN = 60
        img_backbone = models.resnet50(pretrained=True)
        self.img_backbone = nn.Sequential(*list(img_backbone.children())[:-1])
        #print(self.img_backbone)

        self.image_transform = torchvision.transforms.Compose(
    [
        torchvision.transforms.Resize(size=(224, 224)),
        torchvision.transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    def __len__(self):
        #return int(self.csv_data.shape[0]/1000)
        return self.csv_data.shape[0]
    
    def pre_processing_BERT(self, sent):

        # 创建空列表储存输出
        input_ids = []
        attention_mask = []
        
        encoded_sent = self.tokenizer_bert.encode_plus(
            text=text_preprocessing(sent),  # 预处理
            add_special_tokens=True,        # `[CLS]`&`[SEP]`
            max_length=self.MAX_LEN,        # 截断/填充的最大长度
            padding='max_length',           # 句子填充最大长度
            # return_tensors='pt',          # 返回tensor
            return_attention_mask=True,     # 返回attention mask
            truncation=True
            )
        
        input_ids = encoded_sent.get('input_ids')
        attention_mask = encoded_sent.get('attention_mask')
        
        # 转换tensor
        input_ids = torch.tensor(input_ids)
        attention_mask = torch.tensor(attention_mask)
        
        
        return input_ids, attention_mask
     
        
    def __getitem__(self, idx):

        if torch.is_tensor(idx):
            idx = idx.tolist()
        
        img_name = os.path.join(self.root_dir, self.csv_data.iloc[idx]['#2 ImageID'])

        image = Image.open(img_name)
        if image.mode!="RGB":
            image = image.convert("RGB")

        #print(img_name,image.size)

        images = []
        image = self.image_transform(image)
        step_i = 56
        for x in range(4):
            for y in range(4):
                images.append(image[:,x*step_i:(x+1)*step_i,y*step_i:(y+1)*step_i].unsqueeze(dim = 0))

        images = torch.cat(images, dim = 0)
        with torch.no_grad():
            image:torch.Tensor = self.img_backbone(images).squeeze()
        image = image.permute(1,0).reshape([2048,4,4])
        #print(image.permute(1,0).reshape([2048,4,4]).shape)

        text = self.csv_data.iloc[idx]['#3 String']
        tensor_input_id, tensor_input_mask = self.pre_processing_BERT(text)

        label = self.csv_data.iloc[idx]['#1 Label']

        
        label = torch.tensor(label)

        return idx, label, image, [tensor_input_id, tensor_input_mask]
